
#include<src/pcps.h>
#include<src/davidson.h>
#include<src/random.h>
#include<src/blas_lapack.h>

template<class T> class MatrixMultiplier : public pcps::MatrixActor<T> {

  private:

    std::vector<T> _A;

  public:

    MatrixMultiplier(const std::vector<T> & A) : _A(A) {}

    void operate(const int n, const T * const x, T * const y) {
      for (int i = 0; i < n; i++) {
        y[i] = pcps::zero<T>();
        for (int j = 0; j < n; j++)
          y[i] += _A[i*n+j] * x[j];
      }
    }

};

template<class T> class Preconditioner : public pcps::MatrixActor<T> {

  private:

    std::vector<T> _H_diag;
    std::vector<T> _S_diag;
    const double & _eval;

  public:

    Preconditioner(const int n,
                   const std::vector<T> & H,
                   const std::vector<T> & S,
                   const double & eval)
      : _H_diag(n),
        _S_diag(n),
        _eval(eval)
    {
      for (int i = 0; i < n; i++) _H_diag[i] = H[i*n+i];
      for (int i = 0; i < n; i++) _S_diag[i] = S[i*n+i];
    }

    void operate(const int n, const T * const x, T * const y) {
      for (int i = 0; i < n; i++)
        y[i] = x[i] / ( _H_diag[i] - _eval * _S_diag[i] );
    }

};

template<class T> void pcps::test_davidson(const int n) {

  std::vector<T> H(n*n);
  for (int i = 0; i < n; i++)
  for (int j = 0; j < n; j++) {
    H[i*n+j] = 0.1 * pcps::random_number<T>() - 0.05*(pcps::unity<T>() + pcps::imaginary_unity<T>());
    H[j*n+i] = pcps::conj(H[i*n+j]);
  }
  for (int i = 0; i < n; i++)
    H[i*n+i] = -10.0 * pcps::unity<T>() * ( pcps::random_number<double>() + 1.0 );

  std::vector<T> S(n*n);
  for (int i = 0; i < n; i++)
  for (int j = 0; j < n; j++) {
    S[i*n+j] = 0.01 * pcps::random_number<T>() - 0.005*(pcps::unity<T>() + pcps::imaginary_unity<T>());
    S[j*n+i] = pcps::conj(S[i*n+j]);
  }
  for (int i = 0; i < n; i++)
    S[i*n+i] = pcps::unity<T>() * ( 10.0 + 5.0 * pcps::random_number<double>() );

  double E;

  MatrixMultiplier<T> Hop(H);
  MatrixMultiplier<T> Sop(S);
  Preconditioner<T> P(n, H, S, E);

  std::vector<T> evec(n);
  for (int i = 0; i < n; i++)
    evec[i] = pcps::random_number<T>() - 0.5*(pcps::unity<T>() + pcps::imaginary_unity<T>());

  pcps::davidson(n, n+1, 1.0e-10, E, &evec[0], &Hop, &Sop, &P);

  std::vector<T> e1(n, 0.0);
  Hop.operate(n, &evec[0], &e1[0]);

  std::vector<T> e2(n, 0.0);
  Sop.operate(n, &evec[0], &e2[0]);

  std::cout << "final residual vector:" << std::endl;
  for (int i = 0; i < n; i++) {
    e1[i] -= E * e2[i];
    std::cout << pcps::format_number("%20.12f", e1[i]) << std::endl;
  }
  std::cout << std::endl;

}

template void pcps::test_davidson<double>(const int);
template void pcps::test_davidson<std::complex<double> >(const int);
